"""
GARCH波动率交易策略(Zipline实现)
策略逻辑：
1. 使用GARCH模型预测波动率
2. 在低波动时建仓，高波动时减仓
3. 动态调整仓位大小
"""

from zipline.api import order_target_percent, record, symbol, get_datetime
from zipline.finance import commission, slippage
import numpy as np
from arch import arch_model

def initialize(context):
    # 策略参数
    context.asset = symbol('SPY')
    context.lookback = 252  # 使用1年数据
    context.model = None
    context.returns_history = []
    
    # 设置交易成本
    context.set_commission(commission.PerShare(cost=0.001, min_trade_cost=1))
    context.set_slippage(slippage.FixedSlippage(spread=0.01))

def forecast_volatility(context, returns):
    """预测未来波动率"""
    if context.model is None or len(returns) < context.lookback:
        return None
        
    # 拟合GARCH(1,1)模型
    context.model = arch_model(returns, vol='Garch', p=1, q=1)
    results = context.model.fit(disp='off')
    
    # 预测波动率
    forecasts = results.forecast(horizon=5)
    return np.sqrt(forecasts.variance.values[-1][-1])

def handle_data(context, data):
    # 获取当前价格
    current_price = data.current(context.asset, 'price')
    
    # 计算当日收益率
    if len(context.returns_history) > 0:
        prev_price = context.returns_history[-1][0]
        daily_return = (current_price - prev_price) / prev_price
    else:
        daily_return = 0
    
    # 更新收益率历史
    context.returns_history.append((current_price, daily_return))
    
    # 保持历史数据长度
    if len(context.returns_history) > context.lookback:
        context.returns_history.pop(0)
    
    # 检查是否有足够数据
    if len(context.returns_history) < context.lookback:
        return
    
    # 提取收益率序列
    returns_series = np.array([r[1] for r in context.returns_history])
    
    # 预测波动率
    vol = forecast_volatility(context, returns_series)
    
    if vol is None:
        return
    
    # 波动率策略逻辑
    if vol < 0.2:  # 低波动市场
        target_position = 1.0  # 满仓
    elif vol > 0.4:  # 高波动市场
        target_position = -1.0  # 空仓
    else:  # 中等波动
        target_position = 0.5  # 半仓
    
    # 调整仓位
    current_position = context.portfolio.positions[context.asset].amount
    current_value = current_position * current_price if current_position else 0
    current_percent = current_value / context.portfolio.portfolio_value
    
    if abs(target_position - current_percent) > 0.05:
        order_target_percent(context.asset, target_position)
    
    # 记录状态
    record(price=current_price, 
          volatility=vol, 
          position=target_position)